#!/usr/bin/env python
import os
import sys
import math
import random
import itertools
from pyrna import automaton
from pyrna.ahoCorasickAutomaton import RNADesignAutomaton,isForbidden
from optparse import OptionParser, TitledHelpFormatter


def parse(str):
  result = [-1 for i in range(len(str))]
  try:
    tmp = []
    i = 0
    for i in range(len(str)):
      if str[i] == "(":
        tmp.append(i)
      elif str[i] == ")":
        j = tmp.pop()
        result[j] = i
        result[i] = j
  except Exception:
    result = None
  if len(tmp)!=0:
    result = None
  return result

def orderSubIntervals(bp):
  n = len(bp)
  result = []
  intervals = [(0,n-1)]
  while len(intervals)>0:
    (i,j) = intervals.pop()
    if (j>=i):
      result.append((j-i+1,i,j))
      if bp[i]==-1:
        intervals.append((i+1,j))
      else:
        k = bp[i]
        intervals.append((i+1,k-1))
        intervals.append((k+1,j))
  result.sort()
  return [(i,j) for (m,i,j) in result]
      

def nextState(state,trans):
  t = state.getOutTransition(trans)
  if len(t)!=1:
    raise Exception("Not a DFA: Not exactly one transition for [%s] over [%s]\nTransitions: %s"%(state.getName(),trans,state))
  q = t[0].getSucc()
  return q

BASES = {
  'A':['A'],
  'C':['C'],
  'G':['G'],
  'U':['U'],
  'N':['A','C','G','U'],
  'R':['A','G'],
  'Y':['C','U'],
  'S':['G','C'],
  'W':['A','U'],
  'K':['G','U'],
  'M':['A','C'],
  'B':['C','G','U'],
  'D':['A','G','U'],
  'H':['A','C','U'],
  'V':['A','C','G']
  }

ALL_BASES = BASES['N']
ALL_BASES_I = {}
for x in range(len(ALL_BASES)):
  ALL_BASES_I[ALL_BASES[x]] = x

BASES_PAIRS = [('A','U'),('U','A'),('G','C'),('C','G'),('G','U'),('U','G')]

ANY_PAIR = [('A','A'),
            ('A','U'),
            ('A','C'),
            ('A','G'),
            ('U','A'),
            ('U','U'),
            ('U','C'),
            ('U','G'),
            ('C','A'),
            ('C','U'),
            ('C','C'),
            ('C','G'),
            ('G','A'),
            ('G','U'),
            ('G','C'),
            ('G','G'),
            ]

def getBases(c):
  return BASES[c.upper()]

def getBasePairs(c1="N",c2="N"):
  result = list(set(BASES_PAIRS) & set(list(itertools.product(BASES[c1.upper()],BASES[c2.upper()]))))
  return result

RT = 0.61626805

#Ripped from IncaRNAtion
global STACKING_ENERGY
STACKING_ENERGY = {}
for a,b in ANY_PAIR:
  for c,d in ANY_PAIR:
    STACKING_ENERGY[a,b,c,d] = sys.maxsize
#The order of the nucleotides is from 5' -> 3'
STACKING_ENERGY.update({('A', 'A', 'U', 'U'):-0.9,
                        ('A', 'C', 'G', 'U'):-2.2,
                        ('A', 'G', 'C', 'U'):-2.1,
                        ('A', 'G', 'U', 'U'):-0.6,
                        ('A', 'U', 'A', 'U'):-1.1,
                        ('A', 'U', 'G', 'U'):-1.4,
                        ('C', 'A', 'U', 'G'):-2.1,
                        ('C', 'G', 'U', 'G'):-1.4,
                        ('C', 'U', 'A', 'G'):-2.1,
                        ('C', 'U', 'G', 'G'):-2.1,
                        ('G', 'A', 'U', 'C'):-2.4,
                        ('G', 'C', 'G', 'C'):-3.4,
                        ('G', 'G', 'C', 'C'):-3.3,
                        ('G', 'G', 'U', 'C'):-1.5,
                        ('G', 'U', 'A', 'C'):-2.2,
                        ('G', 'U', 'G', 'C'):-2.5,
                        ('G', 'A', 'U', 'U'):-1.3,
                        ('G', 'C', 'G', 'U'):-2.5,
                        ('G', 'G', 'C', 'U'):-2.1,
                        ('G', 'G', 'U', 'U'):-0.5,
                        ('G', 'U', 'A', 'U'):-1.4,
                        ('G', 'U', 'G', 'U'):1.3,
                        ('U', 'A', 'U', 'A'):-1.3,
                        ('U', 'C', 'G', 'A'):-2.4,
                        ('U', 'G', 'C', 'A'):-2.1,
                        ('U', 'G', 'U', 'A'):-1.0,
                        ('U', 'U', 'A', 'A'):-0.9,
                        ('U', 'U', 'G', 'A'):-1.3,
                        ('U', 'A', 'U', 'G'):-1.0,
                        ('U', 'C', 'G', 'G'):-1.5,
                        ('U', 'G', 'C', 'G'):-1.4,
                        ('U', 'G', 'U', 'G'):0.3,
                        ('U', 'U', 'A', 'G'):-0.6,
                        ('U', 'U', 'G', 'G'):-0.5,
                        ('C', 'C', 'G', 'G'):-3.3,
                        ('C', 'G', 'C', 'G'):-2.4})


WEIGHTS = {('A','U'):math.exp(2.0/RT),
 ('U','A'):math.exp(2.0/RT),
 ('G','C'):math.exp(3.0/RT),
 ('C','G'):math.exp(3.0/RT),
 ('G','U'):math.exp(1.0/RT),
 ('U','G'):math.exp(1.0/RT)}

def boltz(e):
  return math.exp(-e/RT)

def weight(bc1,bc2,c1,c2,useWeights,isStacking):
  if useWeights and isStacking:
    return boltz(STACKING_ENERGY[bc1,c1,c2,bc2])
  return 1

def product(l1,l2,l3=[""],l4=[""]):
  result = []
  for (a,b,c,d) in itertools.product(l1,l2,l3,l4):
    result.append(a+b+c+d)
  return result

def countAux(bp,aut,s2i,constraints,listWords=False,useWeights=False):
  n = len(bp)
  intervals = orderSubIntervals(bp)
  states = aut.getStates()
  numStates = len(states)
  mat = [[[[[0.0 for k in range(n)] for j in range(numStates)] for i in range(numStates)] for c1 in ALL_BASES] for c2 in ALL_BASES]
  words = None
  if listWords:
    words = [[[[] for k in range(n)] for j in range(numStates)] for i in range(numStates)]
  numComputations =  len(intervals)*len(states)*len(states)
  numDone = 0
  for (i,j) in intervals:
    for qin in states:    
      for qout in states:
        #if ((numDone%min(1000,math.ceil(numComputations/100.0)))==0):
        #  sys.stderr.write("\r%.2f%%    "%(100.0*numDone/numComputations))
        if bp[i]==-1:
          baseList = getBases(constraints[i])
          for (cb1,cb2) in ANY_PAIR:
            icb1,icb2 = ALL_BASES_I[cb1],ALL_BASES_I[cb2]
            for c in baseList:
              ic = ALL_BASES_I[c]
              qafter = nextState(qin,c)
              if not isForbidden(qafter) or len(baseList)==1:
                if (i+1>j):
                  if qafter == qout:
                    rest = 1.0
                    if listWords: restw = [""]
                  else:
                    rest = 0.0
                    if listWords: restw = []
                else:
                  rest = mat[ic][icb2][s2i[qafter]][s2i[qout]][i+1]
                  if listWords: restw = words[s2i[qafter]][s2i[qout]][i+1]
                mat[icb1][icb2][s2i[qin]][s2i[qout]][i] += rest
                if listWords and ((icb1,icb2)==(0,0)): words[s2i[qin]][s2i[qout]][i] += product([c],restw)
        else:
          k = bp[i]
          isStacking = False
          if (i-1>=0):
            isStacking = (bp[i-1] == k+1)
          pairList = getBasePairs(constraints[i],constraints[k])
          for (cb1,cb2) in ANY_PAIR:
            icb1,icb2 = ALL_BASES_I[cb1],ALL_BASES_I[cb2]
            for (c1,c2) in pairList:
              ic1,ic2 = ALL_BASES_I[c1],ALL_BASES_I[c2]
              #if not isForbidden(qafter1):
              for qmid in states:
                qafter1 = nextState(qin,c1)
                qafter2 = nextState(qmid,c2)
                if ((not isForbidden(qafter1)) and (not isForbidden(qafter2))) or len(pairList)==1:
                  if (i+1>k-1):
                    if qafter1==qmid:
                      rest1 = 1.0
                      if listWords: rest1w = [""]
                    else:
                      rest1 = 0.0
                      if listWords: rest1w = []
                  else:
                    rest1 = mat[ic1][ic2][s2i[qafter1]][s2i[qmid]][i+1]
                    if listWords: rest1w = words[s2i[qafter1]][s2i[qmid]][i+1]
                  if (k+1>j):
                    if qafter2==qout:
                      rest2 = 1.0
                      if listWords: rest2w = [""]
                    else:
                      rest2 = 0.0
                      if listWords: rest2w = []
                  else:
                    rest2 = mat[ic2][icb2][s2i[qafter2]][s2i[qout]][k+1]
                    if listWords: rest2w = words[s2i[qafter2]][s2i[qout]][k+1]
                    
                  mat[icb1][icb2][s2i[qin]][s2i[qout]][i] += weight(cb1,cb2,c1,c2,useWeights,isStacking)*rest1*rest2
                  if listWords and ((icb1,icb2)==(0,0)): words[s2i[qin]][s2i[qout]][i] += product([c1],rest1w,[c2],rest2w)
        numDone += 1
  print(numDone)
  return (mat,words)

def count(bp,aut,constraints=None,listWords=False,useWeights=False):
  s2i = {}
  states = aut.getStates()
  for i in range(len(states)):
    s2i[states[i]] = i
  (mat,words) = countAux(bp,aut,s2i,constraints,listWords,useWeights)
  acc = 0
  for qin in aut.getInitials():    
    for qout in aut.getFinals():
      acc += mat[0][0][s2i[qin]][s2i[qout]][0]
  if listWords:
    for qin in aut.getInitials():    
      for qout in aut.getFinals():
        for w in words[s2i[qin]][s2i[qout]][0]:
          print(w)
  return (acc,mat)


def listAllWords(bp,aut,constraints=None):
  result = []
  s2i = {}
  states = aut.getStates()
  for i in range(len(states)):
    s2i[states[i]] = i
  (mat,words) = countAux(bp,aut,s2i,constraints,True,False)
  for qin in aut.getInitials():    
    for qout in aut.getFinals():
      for w in words[s2i[qin]][s2i[qout]][0]:
        result.append(w)
  return result

def randGenAux(bp,constraints,mat,s2i,qin,qout,i,j=None,icb1=0,icb2=0,useWeights=False):
  if i>j: return ""
  r = random.random()*mat[icb1][icb2][s2i[qin]][s2i[qout]][i]
  cb1,cb2 = ALL_BASES[icb1],ALL_BASES[icb2]
  if bp[i]==-1:
    baseList = getBases(constraints[i])
    for c in baseList:
      ic = ALL_BASES_I[c]
      qafter = nextState(qin,c)
      if not isForbidden(qafter) or len(baseList)==1:
        if (i+1>j):
          if qafter == qout:
            rest = 1.0
          else:
            rest = 0.0
        else:
          rest = mat[ic][icb2][s2i[qafter]][s2i[qout]][i+1]
        r -= rest
        if r<0:
          return c+randGenAux(bp,constraints,mat,s2i,qafter,qout,i+1,j,ic,icb2,useWeights)
  else:
    k = bp[i]
    isStacking = False
    if (i-1>=0):
      isStacking = (bp[i-1] == k+1)
    pairList = getBasePairs(constraints[i],constraints[k])
    for (c1,c2) in pairList:
      ic1,ic2 = ALL_BASES_I[c1],ALL_BASES_I[c2]
      for qmid in list(s2i.keys()):
        qafter1 = nextState(qin,c1)
        qafter2 = nextState(qmid,c2)
        if ((not isForbidden(qafter1)) and (not isForbidden(qafter2))) or len(pairList)==1:
          if (i+1>k-1):
            if qafter1==qmid:
              rest1 = 1.0
            else:
              rest1 = 0.0
          else:
            rest1 = mat[ic1][ic2][s2i[qafter1]][s2i[qmid]][i+1]
          if (k+1>j):
            if qafter2==qout:
              rest2 = 1.0
            else:
              rest2 = 0.0
          else:
            rest2 = mat[ic2][icb2][s2i[qafter2]][s2i[qout]][k+1]
          r -= weight(cb1,cb2,c1,c2,useWeights,isStacking)*rest1*rest2
          if r<0:
            return c1+randGenAux(bp,constraints,mat,s2i,qafter1,qmid,i+1,k-1,ic1,ic2,useWeights)+c2+randGenAux(bp,constraints,mat,s2i,qafter2,qout,k+1,j,ic2,icb2,useWeights)
  return None


def design(bps,aut,constraints=None,nb=1,acc=None,mat=None,useWeights=None):
  n = len(bps)
  if (constraints is None) or (len(constraints)==0):
    constraints = ["N" for i in range(n)]
  s2i = {}
  states = aut.getStates()
  for i in range(len(states)):
    s2i[states[i]] = i
  result = []
  if (mat is None) or (acc is None):
    (acc,mat) = count(bps,aut,constraints,False,useWeights)
    
  for x in range(nb):
    r = random.random()*acc
    for qin in aut.getInitials():    
      for qout in aut.getFinals():
        r -= mat[0][0][s2i[qin]][s2i[qout]][0]
        if r<0:
          result.append(randGenAux(bps,constraints,mat,s2i,qin,qout,0,n-1,0,0,useWeights))
          break
      if r<0: break;
  return result

    
def listConversionCallback(option, opt, value, parser):
  setattr(parser.values, option.dest, value.split(','))


def process_command_line(argv):
  if argv is None:
    argv = sys.argv[1:]

  usage = "usage: %prog [options] [-h|-c|-a|-g num]"
  parser = OptionParser(usage,
                        formatter=TitledHelpFormatter(width=80))
  
  parser.add_option("-s", "--struct", dest="expected_struct",
                    help="Sets expected secondary structure", metavar="str")

  parser.add_option("-w", "--seq", dest="sequence_constraints",
                    help="Sets sequence constraints", metavar="str")

  parser.add_option("-c", "--count", action='store_true',
                    dest="count_sequences",default=False,
                    help="Counts and displays compatible sequences")

  parser.add_option("-e", "--weights", action='store_true',
                    dest="use_weights",default=False,
                    help="Uses Boltzmann weight for sampling, based on a stacking-only version of the Turner 2004 model")

  parser.add_option("-a", "--all", action='store_true',
                    dest="list_sequences",default=False,
                    help="Displays all compatible sequences")

  parser.add_option("-g", "--gen", dest="num_sequences", default=-1,
                    help="Draws a random set of compatible sequences",
                    metavar="int")

  parser.add_option("-m", "--mandatory", dest="mandatory_set",type='string',
                  action='callback',
                  callback=listConversionCallback,
                  help="Sets the list of mandatory motifs",
                  default=[])
  
  parser.add_option("-f", "--forbidden", dest="forbidden_set",type='string',
                  action='callback',
                  callback=listConversionCallback,
                  help="Sets the list of mandatory motifs", 
                  default=[])

  """parser.add_option("--debug", action="store_true", dest="debug",
                    help='Turns on debug messages')"""
  
  """parser.add_option("-o", action='store_false', dest="optimize",
                    help='invokes psyco to optimize')"""
  
  """parser.add_option('-h', '--help', action='help',
                    help='Shows this help message and exit.')"""
  
  (options, args) = parser.parse_args(argv)
  gen_mode = ((int(options.num_sequences) != -1)  and not options.count_sequences and not options.list_sequences)
  if argv:
    if options.expected_struct == None and options.sequence_constraints == None:
      parser.error("User should input either a secondary structure or a sequence mask (or both).")
    elif options.sequence_constraints == None:
      options.sequence_constraints = ["N" for i in range(len(options.expected_struct))]
    elif options.expected_struct == None:
      options.expected_struct = ["." for i in range(len(options.sequence_constraints))]
    if len(options.expected_struct) != len(options.sequence_constraints):
      parser.error("Sequence and structure masks should have the same length.")
    if not gen_mode  and not options.count_sequences and not options.list_sequences:
      parser.error("Please select either of the 'counting' or 'generating' mode.")
    if gen_mode and (int(options.num_sequences) <= 0):
      parser.error("Number of generated sequences must be positive.")
  return options, args


def main(options, args):
    struct = options.expected_struct
    aut = RNADesignAutomaton(options.forbidden_set,options.mandatory_set)
    print(("#States:",len(aut.getStates())))
    bps = parse(options.expected_struct)
    nb = int(options.num_sequences)
    useWeights = options.use_weights
    if options.count_sequences:
      (numSeq,mat) = count(bps,aut,options.sequence_constraints)
      print((int(numSeq)))
    elif options.list_sequences:
      (numSeq,mat) = count(bps,aut,options.sequence_constraints,True)
    elif nb>0:
      (numSeq,mat) = count(bps,aut,options.sequence_constraints,False,useWeights)
      print(numSeq)
      for s in design(bps,aut,options.sequence_constraints,nb,numSeq,mat,useWeights):
        print(s)

###############################################################################

if __name__ == '__main__':
    options, args = process_command_line(sys.argv)
    main(options, args)
